import os

from langchain_classic.chat_models import init_chat_model
# from langchain.chat_models import init_chat_model
from pydantic import BaseModel, Field, field_validator
from typing import Literal
from langchain.tools import tool
from langchain_classic.agents import create_tool_calling_agent, AgentExecutor
from langchain_core.prompts import ChatPromptTemplate

from dotenv import load_dotenv
# 加载.env文件
load_dotenv()

# 1. 输入参数模型（支持中文描述）
class WeatherInput(BaseModel):
    """天气查询工具的输入参数，支持中文地址和单位选择"""
    location: str = Field(
        ...,  # 必填项
        description="城市名称（如'北京'、'上海'）或经纬度（如'39.9042° N, 116.4074° E'），支持中文"
    )
    units: Literal["celsius", "fahrenheit", "摄氏度", "华氏度"] = Field(  # 新增中文单位选项
        default="摄氏度",
        description="温度单位：'celsius'/'摄氏度'（°C）或'fahrenheit'/'华氏度'（°F）"
    )
    include_forecast: bool = Field(
        default=False,
        description="是否包含5天预报，True为包含，False仅返回当前天气"
    )

    # 验证地址不为空
    @field_validator('location')
    def location_cannot_be_empty(cls, v):
        if not v.strip():
            raise ValueError("地址不能为空，请输入城市名称或经纬度")
        return v

# 2. 天气工具（兼容中文地址和单位）
@tool(
    "weather_tool",
    args_schema=WeatherInput,
    description="查询指定地点的天气（支持中文地址，如'广州'），用户问天气时必须使用此工具"
)
def get_weather(
    location: str,
    units: str = "摄氏度",
    include_forecast: bool = False
) -> str:
    """获取指定地点的当前天气和可选的5天预报，支持中文地址解析"""
    try:
        # 统一处理单位（将中文单位转换为英文，方便后续逻辑）
        units_map = {"摄氏度": "celsius", "华氏度": "fahrenheit"}
        units = units_map.get(units.lower(), units.lower())  # 兼容中英文单位输入

        # 模拟天气数据库（包含中文城市名）
        weather_db = {
            "北京": {"current_temp": 22, "condition": "晴"},
            "上海": {"current_temp": 25, "condition": "多云"},
            "广州": {"current_temp": 28, "condition": "小雨"},
            "beijing": {"current_temp": 22, "condition": "晴"},  # 兼容英文拼写
            "39.9042° n, 116.4074° e": {"current_temp": 22, "condition": "晴"}  # 经纬度
        }

        # 处理中文地址（统一转为小写，移除空格，提高匹配度）
        location_key = location.strip().lower()
        # 特殊处理：如果是中文城市，保留原文字（因为小写会导致拼音问题）
        if any(char >= '\u4e00' and char <= '\u9fff' for char in location_key):
            location_key = location.strip()

        # 检查地址是否存在
        if location_key not in weather_db:
            return f"抱歉，暂时无法获取'{location}'的天气数据，请尝试其他地点"

        # 获取温度并转换单位
        data = weather_db[location_key]
        current_temp = data["current_temp"]
        if units == "fahrenheit":
            current_temp = current_temp * 9/5 + 32  # 摄氏度转华氏度

        # 构建结果（根据单位显示对应的符号）
        unit_symbol = "°C" if units == "celsius" else "°F"
        result = (
            f"{location}的当前天气：\n"
            f"温度：{current_temp:.1f}{unit_symbol}\n"
            f"天气状况：{data['condition']}"
        )

        # 追加5天预报
        if include_forecast:
            forecast = [
                "第1天：24°C，晴",
                "第2天：23°C，多云",
                "第3天：25°C，晴",
                "第4天：22°C，阴",
                "第5天：20°C，小雨"
            ]
            # 转换预报单位（如果是华氏度）
            if units == "fahrenheit":
                forecast = [
                    day.replace("°C", "°F").replace(
                        str(int(day.split("°")[0])),
                        str(int(int(day.split("°")[0]) * 9/5 + 32))
                    )
                    for day in forecast
                ]
            result += "\n\n5天预报：\n" + "\n".join(forecast)

        return result

    except Exception as e:
        return f"获取天气失败：{str(e)}"

# 3. 集成Agent（支持中文查询）
if __name__ == "__main__":
    model = init_chat_model(
        model='deepseek-reasoner',  # deepseek-chat表示调用DeepSeek-v3模型，deepseek-reasoner表示调用DeepSeek-R1模型，
        model_provider='deepseek',  # 模型提供商写deepseek
        # api_key="sk-476b894aa2d74474b77d4dfa91633f72", #你注册的deepseek api_key
        api_key=os.environ.get("DEEPSEEK_API_KEY"),
    )

    # 提示词明确支持中文地址
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是天气助手，支持中文查询。当用户问天气时，必须调用'weather_tool'，"
                   "地址可以是中文城市名（如'成都'），单位可以用中文（如'摄氏度'）。"
                   "如果用户没说单位，默认用摄氏度。"),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}")
    ])

    agent = create_tool_calling_agent(model, [get_weather], prompt)
    agent_executor = AgentExecutor(
        agent=agent,
        tools=[get_weather],
        verbose=True,
        handle_parsing_errors=True
    )

    # 测试中文地址查询
    test_queries = [
        "北京的天气怎么样？",
        "上海的温度是多少，用华氏度显示",
        "广州的天气，包含未来5天预报"
    ]

    for query in test_queries:
        print(f"\n用户：{query}")
        response = agent_executor.invoke({"input": query})
        print(f"助手：{response['output']}")